//////////////////////////////////
//
// 这个类完成的是一个图像畸变矫正的工作
// 使用的是针孔相机模型
//
/////////////////////////////////
#include "vision_ruler/distortion_correction.h"

vision_ruler::Distortion_Correction_Util::Distortion_Correction_Util() {}
vision_ruler::Distortion_Correction_Util::~Distortion_Correction_Util() {}

void vision_ruler::Distortion_Correction_Util::doCorrectPoint(std::vector<vision_ruler::box_info>& boxes, vision_ruler::camera_info cinfo) {
    for(size_t i = 0; i < boxes.size(); i++) {
        vision_ruler::box_info & box = boxes[i];
        cv::Point2f left_top_pixel_corrected = imageCorrection(box.left_top_pixel, cinfo);
        cv::Point2f right_bottom_pixel_corrected = imageCorrection(box.right_bottom_pixel, cinfo);
        int outOfRangeNum = 0;
        int min_x, max_x, min_y, max_y;
        double ltx, lty, rbx, rby;
        ltx = left_top_pixel_corrected.x;
        lty = left_top_pixel_corrected.y;
        rbx = right_bottom_pixel_corrected.x;
        rby = right_bottom_pixel_corrected.y;
        if(left_top_pixel_corrected.x < 0 || left_top_pixel_corrected.y < 0 || left_top_pixel_corrected.x >= cinfo.width || left_top_pixel_corrected.y >= cinfo.height) {
            if(left_top_pixel_corrected.x < 0) left_top_pixel_corrected.x = 0;
            if(left_top_pixel_corrected.x >= cinfo.width) left_top_pixel_corrected.x = cinfo.width - 1;
            if(left_top_pixel_corrected.y < 0) left_top_pixel_corrected.y = 0;
            if(left_top_pixel_corrected.y >= cinfo.height) left_top_pixel_corrected.y = cinfo.height - 1;
            outOfRangeNum++;
        }
        if(right_bottom_pixel_corrected.x < 0 || right_bottom_pixel_corrected.y < 0 || right_bottom_pixel_corrected.x >= cinfo.width || right_bottom_pixel_corrected.y >= cinfo.height) {
            if(right_bottom_pixel_corrected.x < 0) right_bottom_pixel_corrected.x = 0;
            if(right_bottom_pixel_corrected.x >= cinfo.width) right_bottom_pixel_corrected.x = cinfo.width - 1;
            if(right_bottom_pixel_corrected.y < 0) right_bottom_pixel_corrected.y = 0;
            if(right_bottom_pixel_corrected.y >= cinfo.height) right_bottom_pixel_corrected.y = cinfo.height - 1;
            outOfRangeNum++;
        }
        if(outOfRangeNum == 2) {
            box.isValid = false;
            box.distance = -1;
            continue;
        }
        box.widthScale = (right_bottom_pixel_corrected.x - left_top_pixel_corrected.x) / (rbx - ltx);
        box.heightScale = (right_bottom_pixel_corrected.y - left_top_pixel_corrected.y) / (rby - lty);
        box.left_top_pixel.x = left_top_pixel_corrected.x;
        box.left_top_pixel.y = left_top_pixel_corrected.y;
        box.right_bottom_pixel.x = right_bottom_pixel_corrected.x;
        box.right_bottom_pixel.y = right_bottom_pixel_corrected.y;
    }
}

cv::Point2f vision_ruler::Distortion_Correction_Util::imageCorrection(cv::Point2f pt, vision_ruler::camera_info cinfo) {
    // 针孔相机 + 径向畸变 + 切向畸变
    double k1 = cinfo.k1, k2 = cinfo.k2, k3 = cinfo.k3, p1 = cinfo.p1, p2 = cinfo.p2;
    double fx = cinfo.fx, fy = cinfo.fy, cx = cinfo.cx, cy = cinfo.cy;
    double u, v, x, y, r2, xy, x2, y2, xcorr_rad, ycorr_rad, xcorr_tan, ycorr_tan, xcorr, ycorr, ucorr, vcorr;
    u = pt.x;
    v = pt.y;
    x = (u - cx) / fx;
    y = (v - cy) / fy;
    r2 = pow(x, 2) + pow(y, 2);
    xy = x * y;
    x2 = x * x;
    y2 = y * y;
    xcorr_rad = x * (k1 * r2 + k2 * r2 * r2 + k3 * r2 * r2 * r2);
    ycorr_rad = y * (k1 * r2 + k2 * r2 * r2 + k3 * r2 * r2 * r2);
    xcorr_tan = 2 * p1 * xy + p2 * (r2 + 2 * x2);
    ycorr_tan = p1 * (r2 + 2 * y2) + 2 * p2 * xy;
    xcorr = x + xcorr_rad + xcorr_tan;
    ycorr = y + ycorr_rad + ycorr_tan;
    ucorr = xcorr * fx + cx;
    vcorr = ycorr * fy + cy;
    cv::Point2f _pt(ucorr, vcorr);
    return _pt;
}
